# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import math

import coin.proto.coin_order_enums_pb2 as coin_order_pb2

from coin.proto.coin_query_pb2 import ProductFillElement
from coin.exchange.base.order_gateway import OrderSide
from coin.exchange.base.order_gateway_info import BaseOrderGatewayInfoMutator
from coin.proto.coin_query_pb2 import (
    AccountBalance,
    AccountPosition,
)


def from_proto_order_side(order_side):
  order_side_map = {
      coin_order_pb2.BUY_ORDER: OrderSide.BUY,
      coin_order_pb2.SELL_ORDER: OrderSide.SELL,
      coin_order_pb2.BUY_OPEN_ORDER: OrderSide.BUY,
      coin_order_pb2.BUY_CLOSE_ORDER: OrderSide.BUY,
      coin_order_pb2.SELL_OPEN_ORDER: OrderSide.SELL,
      coin_order_pb2.SELL_CLOSE_ORDER: OrderSide.SELL,
  }
  return order_side_map[order_side]


def _order_side_to_sign(order_side: OrderSide):
  assert isinstance(order_side, OrderSide)
  if order_side == OrderSide.BUY:
    sign = +1  # increase base, decrease quote, or increase net position
  elif order_side == OrderSide.SELL:
    sign = -1  # decrease quote, increase base, or decrease net position
  else:
    raise ValueError('Invalid order side! %s' % order_side)
  return sign


class BaseBalanceOrderGatewayInfo(BaseOrderGatewayInfoMutator):
  @property
  def exchange(self):
    # noinspection PyProtectedMember
    return self.CurrencyType._exchange

  def update_by_fill(self, fill_proto: ProductFillElement):
    product = self.ProductType.FromStr(fill_proto.symbol)
    fee_currency = None
    if fill_proto.HasField('fee_currency'):
      fee_currency = self.CurrencyType.FromStr(fill_proto.fee_currency)
    order_side = from_proto_order_side(fill_proto.side)

    sign = _order_side_to_sign(order_side)
    total = self.get_balance(product.base)
    total += (sign * fill_proto.qty)
    if fee_currency and fee_currency == product.base:
      total -= fill_proto.fee
    self.set_balance(product.base, total)

    total = self.get_balance(product.quote)
    total -= (sign * fill_proto.qty * fill_proto.price)
    if fee_currency and fee_currency == product.quote:
      total -= fill_proto.fee
    self.set_balance(product.quote, total)

  def balance_as_proto(self):
    return AccountBalance(exchange=self.exchange, each_balance=self.get_notNone_balances_internal())

  def isclose(self, og_info, *, rel_tol, abs_tol=1e-4):
    assert 0.0 < rel_tol < 0.1, rel_tol
    assert 0.0 < abs_tol < 0.1, abs_tol
    assert self.currencies == og_info.currencies, (self.currencies, og_info.currencies)
    diff = {}
    for currency in self.currencies:
      diff[currency.currency] = math.isclose(self.get_balance(currency),
                                             (og_info.get_balance(currency) or 0),
                                             rel_tol=rel_tol,
                                             abs_tol=abs_tol)
    return all(diff.values())

  def clone(self):
    return type(self)(self._currencies, self._products)


class BasePositionOrderGatewayInfo(BaseOrderGatewayInfoMutator):
  @property
  def exchange(self):
    # noinspection PyProtectedMember
    return self.CurrencyType._exchange

  def update_by_fill(self, fill_proto: ProductFillElement):
    product = self.ProductType.FromStr(fill_proto.symbol)
    order_side = from_proto_order_side(fill_proto.side)

    sign = _order_side_to_sign(order_side)
    total = self.get_position(product)
    total += (sign * fill_proto.qty)
    self.set_position(product, total)

  def balance_as_proto(self):
    return AccountBalance(exchange=self.exchange, each_balance=self.get_notNone_balances_internal())

  def position_as_proto(self):
    return AccountPosition(exchange=self.exchange,
                           each_position=self.get_notNone_positions_internal())

  def isclose(self, og_info, *, rel_tol, abs_tol=1e-4):
    assert 0.0 < rel_tol < 0.1, rel_tol
    assert 0.0 < abs_tol < 0.1, abs_tol
    assert self.products == og_info.products, (self.products, og_info.products)
    diff = {}
    for product in self.products:
      diff[product] = math.isclose(self.get_position(product), (og_info.get_position(product) or 0),
                                   rel_tol=rel_tol,
                                   abs_tol=abs_tol)
    return all(diff.values())

  def get_product_position_proto(self, product):
    return self._positions_internal[product.symbol]

  def get_average_entry_price(self, product):
    return self._positions_internal[product.symbol].avg_entry_price

  def clone(self):
    return type(self)(self._currencies, self._products)
